package cool.taomu.toolkit.executor

import com.google.common.base.Function
import com.google.common.util.concurrent.AsyncFunction
import com.google.common.util.concurrent.FutureCallback
import com.google.common.util.concurrent.Futures
import com.google.common.util.concurrent.ListenableFuture
import java.util.concurrent.Executor
import java.util.concurrent.ScheduledExecutorService
import java.util.concurrent.TimeUnit
import java.util.function.Consumer
import java.util.function.LongConsumer
import java.util.function.Predicate
import java.util.function.Supplier
import org.slf4j.LoggerFactory

class FutureChainBuilder<T> {

    protected ListenableFuture<T> future;
    protected final Executor executor;
    protected final ScheduledExecutorService scheduler;

    val static LOG = LoggerFactory.getLogger(FutureChainBuilder)

    new(ListenableFuture<T> future, Executor executor, ScheduledExecutorService scheduler) {
        this.future = future;
        this.executor = executor;
        this.scheduler = scheduler;
    }

    def build() {
        this.future
    }

    def onSuccess(Consumer<T> onSuccess) {
        Futures.addCallback(future, new FutureCallback<T>() {
            override void onSuccess(T result) { onSuccess.accept(result); }

            override void onFailure(Throwable t) {}
        }, executor);
        return this
    }

    def onFailure(Consumer<Throwable> onFailure) {
        Futures.addCallback(future, new FutureCallback<T>() {
            override void onSuccess(T result) {}

            override void onFailure(Throwable t) { onFailure.accept(t); }
        }, executor);
        return this
    }

    def callback(FutureCallback<T> fcallback) {
        Futures.addCallback(future, fcallback, executor);
        return this
    }

    def log(String message) {
        Futures.addCallback(future, new FutureCallback<T>() {
            override void onSuccess(T result) {
                LOG.info(message + " - Success: " + result);
            }

            override void onFailure(Throwable t) {
                LOG.error(message + " - Failure: " + t.getMessage());
            }
        }, executor);
        return this
    }

    def timeIt(LongConsumer timeRecorder) {
        val long startTime = System.currentTimeMillis();
        future = Futures.transform(future, [ result |
            var long duration = System.currentTimeMillis() - startTime;
            timeRecorder.accept(duration);
            return result;
        ], executor);
        return this
    }

    def retry(int maxRetries) {
        var ListenableFuture<T> retryFuture = future;
        for (var int i = 0; i < maxRetries; i++) {
            val ListenableFuture<T> finalRetryFuture = retryFuture;
            retryFuture = Futures.catchingAsync(finalRetryFuture, Throwable, [finalRetryFuture], executor);
        }
        future = retryFuture;
        return this
    }

    def retry(int maxRetries, long initialDelay, TimeUnit unit) {
        var ListenableFuture<T> retryFuture = future;
        for (var int attempt = 1; attempt <= maxRetries; attempt++) {
            val long delay = (initialDelay * Math.pow(2, attempt - 1)) as long;
            val ListenableFuture<T> finalRetryFuture = retryFuture;
            retryFuture = Futures.catchingAsync(finalRetryFuture, Throwable, [ ex |
                return Futures.withTimeout(Futures.immediateFailedFuture(ex), delay, unit, scheduler);
            ], executor);
        }
        future = retryFuture;
        return this
    }

    def filter(Predicate<T> predicate, Supplier<Throwable> exceptionSupplier) {
        future = Futures.transform(future, [ result |
            if (!predicate.test(result)) {
                throw exceptionSupplier.get();
            }
            return result;
        ], executor);
        return this
    }

    def withMetrics(String stageName) {
        Futures.addCallback(future, new FutureCallback<T>() {
            override void onSuccess(T result) {
                System.out.println("Metrics: " + stageName + " success");
            }

            override void onFailure(Throwable t) {
                System.out.println("Metrics: " + stageName + " failure");
            }
        }, executor);
        return this
    }

    def catching(Function<Throwable, T> fallback) {
        future = Futures.catching(future, Throwable, fallback, executor);
        return this
    }

    def catchingAsync(Function<Throwable, ListenableFuture<T>> fallback) {
        future = Futures.catchingAsync(future, Throwable, [ex|fallback.apply(ex)], executor);
        return this
    }

    def withTimeout(long timeout, TimeUnit unit) {
        future = Futures.withTimeout(future, timeout, unit, scheduler);
        return this
    }

    def transform(Function<? super T, ? extends T> function) {
        future = Futures.transform(future, function, executor);
        return this
    }

    def flatMap(AsyncFunction<? super T, ? extends T> function) {
        future = Futures.transformAsync(future, function, executor);
        return this
    }

}
